import numpy as np
import matplotlib.pyplot as plt
import os
from matplotlib.colors import LogNorm

fname = 'exfig8b.npz'

# load npz file
data = np.load(fname, allow_pickle=True)
print(data.files)
B_lin = data['B_lin']
Rxy_asym = data['Rxy_asym']
Rxx_sym = data['Rxx_sym']
Rxy_down = data['Rxy_down']
Rxy_up = data['Rxy_up']
Rxx_down = data['Rxx_down']
Rxx_up = data['Rxx_up']
indices = data['indices']
T = data['T']
nu = data['nu']
n = data['n']
D = data['D']
Vtg = data['Vtg']
Vbg = data['Vbg']
n0 = data['n0'] # density at charge neutrality point
D0 = data['D0'] 
ns = data['ns'] # density at 4 electron per moire unit cell
device_name = data['device_name']
back_gate_thickness_nm = data['back_gate_thickness_nm']
top_gate_thickness_nm = data['top_gate_thickness_nm']

# plot
fig = plt.figure(figsize=(8,9))
fig.suptitle('exfig8b')
ax = fig.add_subplot(111)
waterfall_xy = -0.4

lines = []
colors = []
strings = []
# indices = range(0,int(T3.ny/2),1)
indices = list(range(0,72,2))
nu_max = max(nu[indices])
nu_min = min(nu[indices])

for m, index in enumerate(indices):
    # color = plt.cm.Spectral_r((nu3[index]-nu_min)/(nu_max-nu_min))
    color = plt.cm.plasma_r((nu[index]-nu_min)/(nu_max-nu_min))
    colors.append(color)
    string = "{:.2f}".format(nu[index])
    print('m,index,string',m,index,string)
    strings.append(string)
    l, = ax.plot(B_lin, Rxy_asym[:,index//2]/1e3 + waterfall_xy*m,color=color,label=string)
    lines.append(l)
    l, = ax.plot(B_lin, -Rxy_asym[::-1,index//2]/1e3 + waterfall_xy*m,color=color,linestyle='--')
    lines.append(l)

ax.set_xlabel('$B$ (T)')
ax.set_ylabel('$R_{yx}$ (k$\Omega$)')
plt.legend()
plt.show()